import os
import shutil
from concurrent.futures import ThreadPoolExecutor, as_completed

# 源路径
source_base = r"E:\zhihao\2025文档\25年10月\ai_训练\ai项目_备份\yangzhihao\ai_project\public_security\traffic_accient\train\traffic_accidents_model\datasets"
source_train = os.path.join(source_base, "train")
source_valid = os.path.join(source_base, "valid")

# 目标路径
dest_base = r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\交通事故数据集_重清洗_20251022\gta_yx_数据"
dest_train = os.path.join(dest_base, "train")
dest_valid = os.path.join(dest_base, "valid")

# 创建目标文件夹结构
dest_folders = [
    os.path.join(dest_train, "images"),
    os.path.join(dest_train, "labels"),
    os.path.join(dest_valid, "images"),
    os.path.join(dest_valid, "labels")
]

for folder in dest_folders:
    os.makedirs(folder, exist_ok=True)

print("✅ 已创建所有目标文件夹")
print(f"  - {dest_train}\\images")
print(f"  - {dest_train}\\labels")
print(f"  - {dest_valid}\\images")
print(f"  - {dest_valid}\\labels")


def copy_file_pair(filename, src_img_dir, src_lbl_dir, dst_img_dir, dst_lbl_dir):
    """复制一对文件(图片和对应的标签)"""
    try:
        # 复制图片
        src_img = os.path.join(src_img_dir, filename)
        dst_img = os.path.join(dst_img_dir, filename)
        shutil.copy2(src_img, dst_img)

        # 查找并复制对应的标签文件
        base_name = os.path.splitext(filename)[0]
        label_copied = False

        # 尝试常见的标签扩展名
        for ext in ['.txt', '.xml', '.json']:
            label_file = base_name + ext
            src_lbl = os.path.join(src_lbl_dir, label_file)
            if os.path.exists(src_lbl):
                dst_lbl = os.path.join(dst_lbl_dir, label_file)
                shutil.copy2(src_lbl, dst_lbl)
                label_copied = True
                break

        return (True, label_copied)
    except Exception as e:
        print(f"  ⚠️  复制失败: {filename} - {e}")
        return (False, False)


def process_dataset(dataset_type, source_dir, dest_dir):
    """处理单个数据集(train 或 valid)"""
    print(f"\n{'=' * 60}")
    print(f"处理 {dataset_type} 数据集")
    print(f"{'=' * 60}")

    source_images = os.path.join(source_dir, "images")
    source_labels = os.path.join(source_dir, "labels")
    dest_images = os.path.join(dest_dir, "images")
    dest_labels = os.path.join(dest_dir, "labels")

    # 获取所有包含 "yx" 的图片文件
    print("扫描文件...")
    yx_files = []

    if not os.path.exists(source_images):
        print(f"⚠️  源目录不存在: {source_images}")
        return 0, 0

    for filename in os.listdir(source_images):
        file_path = os.path.join(source_images, filename)
        if os.path.isfile(file_path):
            # 只选择包含 "yx" 的文件
            if "yx" in filename.lower():
                yx_files.append(filename)

    print(f"找到 {len(yx_files)} 个包含 'yx' 的图片文件")

    if len(yx_files) == 0:
        print("没有找到包含 'yx' 的文件,跳过...")
        return 0, 0

    # 使用多线程复制文件
    print("开始复制文件...")
    success_count = 0
    labels_found = 0

    with ThreadPoolExecutor(max_workers=8) as executor:
        futures = [
            executor.submit(copy_file_pair, f, source_images, source_labels, dest_images, dest_labels)
            for f in yx_files
        ]

        for i, future in enumerate(as_completed(futures), 1):
            img_ok, lbl_ok = future.result()
            if img_ok:
                success_count += 1
            if lbl_ok:
                labels_found += 1

            if i % 50 == 0 or i == len(yx_files):
                print(f"  进度: {i}/{len(yx_files)}")

    print(f"\n{dataset_type} 统计:")
    print(f"  - 图片: {success_count} 个")
    print(f"  - 标签: {labels_found} 个")

    return success_count, labels_found


# 主程序
print("\n" + "=" * 60)
print("开始复制包含 'yx' 的文件")
print("=" * 60)

# 处理 train 数据集
train_imgs, train_lbls = process_dataset("train", source_train, dest_train)

# 处理 valid 数据集
valid_imgs, valid_lbls = process_dataset("valid", source_valid, dest_valid)

# 总结
print("\n" + "=" * 60)
print("✅ 任务完成!")
print("=" * 60)
print(f"\ntrain 数据集:")
print(f"  - 图片: {train_imgs} 个")
print(f"  - 标签: {train_lbls} 个")
print(f"\nvalid 数据集:")
print(f"  - 图片: {valid_imgs} 个")
print(f"  - 标签: {valid_lbls} 个")
print(f"\n总计:")
print(f"  - 图片: {train_imgs + valid_imgs} 个")
print(f"  - 标签: {train_lbls + valid_lbls} 个")
print("=" * 60)

print(f"\n目标文件夹: {dest_base}")